package geppetto.cat.models;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.EOFException;
import java.io.FileInputStream;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.util.zip.GZIPInputStream;
import java.util.zip.GZIPOutputStream;

//TODO JAVG
//DOES NOT WORK IF BIGGEST SENTENCE IS SMALLER THAN NUMBER OF BINS

/**
 * a distortion table that puts distortions into bins based on the difference
 * between the current source word and the previous source word. Probability
 * mass that extends past the end ot the table is spread uniformly over the rest
 * of the length of the sentence.
 */
public class DistortionTable {

	/** distribution on starting out in state */
	float[] firstProbs;

	/** distribution on starting out in state */
	float[] lastProbs;

	/** unnormalized probabilities */
	float[] distProbs;

	/** unnormalized probability to be split among long sentences */
	float positiveOverflow;

	/** unnormalized probability to be split among long sentences */
	float negativeOverflow;

	/** where zero lies in the distProbs variable */
	int diffZero;

	/*
	 * These are replaced by nullDistProb = 1/(source.length +1) final float
	 * nullDistProb = 0.2f; final float nonNullDistProb = (1-nullDistProb);
	 */

	private Cache cache;

	/**
	 * make a table with the specified number of bins (plus positive and
	 * negative overflow).
	 */
	public DistortionTable(int numBins, int cacheSize) {
		// FIXME debug code to remove
		Runtime.getRuntime().gc();
		double memory = (Runtime.getRuntime().totalMemory() - Runtime
				.getRuntime().freeMemory())
				/ (1024 * 1024);

		// /////

		cache = new Cache(cacheSize);
		distProbs = new float[numBins];
		firstProbs = new float[numBins];
		lastProbs = new float[numBins];
		diffZero = numBins / 2;

		// FIXME debug code to remove
		Runtime.getRuntime().gc();
		double memoryAfter = (Runtime.getRuntime().totalMemory() - Runtime
				.getRuntime().freeMemory())
				/ (1024 * 1024);
		System.out.println("HMM Distortion table  " + (memoryAfter - memory)
				+ " MB");

		// /

	}

	/**
	 * initialize the distortion table with uniform unnormalized probabilities.
	 */
	public DistortionTable(int numBins, float initialConstant, int cacheSize) {
		this(numBins, cacheSize);
		uniform(initialConstant);
	}

	public DistortionTable(String directory) {
		System.out.println("------------Loading HMM DistortionTable");
		String inFilename = directory + "HMMDistortion.gz";
		try {
			DataInputStream data_in = new DataInputStream(new GZIPInputStream(
					new FileInputStream(inFilename)));
			int numBins = data_in.readInt();
			int cacheSize = data_in.readInt();
			cache = new Cache(cacheSize);
			distProbs = new float[numBins];
			firstProbs = new float[numBins];
			lastProbs = new float[numBins];
			diffZero = numBins / 2;
			for (int i = 0; i < numBins; i++) {
				// while (true) {
				try {
					int position = data_in.readInt();
					float prob = data_in.readFloat();
					float first = data_in.readFloat();
					float last = data_in.readFloat();
					distProbs[position] = prob;
					firstProbs[position] = first;
					lastProbs[position] = last;
				} catch (EOFException eof) {
					System.out.println("FINISHED READING FILE");
					break;

				}
			}
			positiveOverflow = data_in.readFloat();
			negativeOverflow = data_in.readFloat();
			data_in.close();
		} catch (FileNotFoundException e) {
			System.out.println("HMM Distortion Table does not exits");
			System.exit(1);
		} catch (IOException e) {
			System.out.println("Error reading HMM Distortion");
			System.exit(1);
		}
		System.out.println("------------Loading Table END");
	}

	public boolean equals(DistortionTable dist) {
		if (distProbs.length != dist.distProbs.length)
			return false;
		for (int i = 0; i < distProbs.length; i++) {
			if (distProbs[i] != dist.distProbs[i])
				return false;
			if (firstProbs[i] != dist.firstProbs[i])
				return false;
			if (lastProbs[i] != dist.lastProbs[i])
				return false;
		}
		return true;
	}

	public void saveTable(String directory) {
		try {
			System.out.println("------------Saving HMM Distortion Table");
			String outFilename = directory + "HMMDistortion.gz";
			DataOutputStream out = new DataOutputStream(new GZIPOutputStream(
					new FileOutputStream(outFilename)));
			out.writeInt(distProbs.length); // Number of bins
			out.writeInt(cache.cacheSize); // CacheSize
			for (int i = 0; i < distProbs.length; i++) {
				out.writeInt(i);
				out.writeFloat(distProbs[i]);
				out.writeFloat(firstProbs[i]);
				out.writeFloat(lastProbs[i]);
			}
			out.writeFloat(positiveOverflow);
			out.writeFloat(negativeOverflow);
			out.close();
		} catch (IOException e) {
			System.out.println("Error saving HMM Distortion Table");
			System.exit(1);
		}
		System.out.println("------------Saving HMM Distortion Table END");
	}

	/** set all entries to 0. */
	public void clear() {
		uniform(0);
	}

	/** set all entries to a uniform value */
	private void uniform(float v) {
		if (v != 0) {
			for (int i = 0; i < firstProbs.length; i++)
				firstProbs[i] = 1f / firstProbs.length;
			for (int i = 0; i < lastProbs.length; i++)
				lastProbs[i] = 1f / lastProbs.length;
		} else {
			for (int i = 0; i < firstProbs.length; i++)
				firstProbs[i] = 0;
			for (int i = 0; i < lastProbs.length; i++)
				lastProbs[i] = 0;
		}
		for (int i = 0; i < distProbs.length; i++)
			distProbs[i] = v;
		positiveOverflow = v;
		negativeOverflow = v;
		cache.invalidate();
	}

	/**
	 * Copy the other distortion table to this one and also divide all the
	 * values by their max
	 */
	public void copy(DistortionTable other) {
		if (distProbs.length != other.distProbs.length)
			throw new IllegalArgumentException("table sizes differ!");
		float max = 0;
		for (int i = 0; i < distProbs.length; i++) {
			if (max < other.distProbs[i])
				max = other.distProbs[i];
		}
		if (max < other.positiveOverflow)
			max = other.positiveOverflow;
		if (max < other.negativeOverflow)
			max = other.negativeOverflow;
		for (int i = 0; i < distProbs.length; i++) {
			distProbs[i] = other.distProbs[i] / max;
		}
		positiveOverflow = other.positiveOverflow / max;
		negativeOverflow = other.negativeOverflow / max;

		float sum = 0;
		for (float f : other.firstProbs)
			sum += f;
		for (int i = 0; i < firstProbs.length; i++)
			firstProbs[i] = other.firstProbs[i] / sum;
		sum = 0;
		for (float f : other.lastProbs)
			sum += f;
		for (int i = 0; i < lastProbs.length; i++)
			lastProbs[i] = other.lastProbs[i] / sum;
		cache.invalidate();
	}

	/**
	 * get the distortion probability given that we are notNull (i.e. don't use
	 * the nullDistProb here).
	 */
	private final float calcGivenNonnull(int si, int sprev, int len) {
		float norm = getNorm(sprev, len);
		assert (norm != 0);
		assert (si < len);
		assert (sprev < len);
		int diff = si - sprev + diffZero;
		if (diff < 0) {
			return tooSmall(sprev, len) / getNorm(sprev, len);
		} else if (diff >= distProbs.length) {
			return tooBig(sprev, len) / getNorm(sprev, len);
		} else {
			return distProbs[diff] / getNorm(sprev, len);
		}
	}

	/**
	 * get the probability that we align to position si given that we are in a
	 * sentence of length len and the previous word aligned to position sprev.
	 * FIXME -- this could be made much more efficient (at least we could save
	 * normalization).
	 */
	public final float getDistProb(int si, int sprev, int len) {

		float nullDistProb = 1f / (1 + len);
		float nonNullDistProb = 1 - nullDistProb;

		assert (getNorm(sprev, len) != 0) : "getting dist norm " + sprev + " "
				+ len;
		assert (si < len);
		assert (sprev < len);

		if (!cache.cachesValid)
			cache.reset();
		if (cache.getVal(si, sprev, len) < 0) {
			cache.setVal(si, sprev, len, nonNullDistProb
					* calcGivenNonnull(si, sprev, len));
		}
		assert !Double.isNaN(cache.getVal(si, sprev, len));
		return cache.getVal(si, sprev, len);

	}

	public final float getLastProb(int si, int len) {
		float nullDistProb = 1f / (1 + len);
		float nonNullDistProb = 1 - nullDistProb;
		float mult = nonNullDistProb;
		if (si >= len) {
			mult = nullDistProb;
			si -= len;
		}
		float norm = 1f;
		if (len < lastProbs.length) {
			float sum = 0f;
			for (int i = 0; i < len; i++) {
				sum += lastProbs[i];
			}
			norm = 1f / sum;
		}
		mult *= norm;

		// we index the lastProb array with len-si, since we care about
		// modeling the end well.
		si = (len - 1) - si;
		if (si < lastProbs.length - 1)
			return mult * lastProbs[si];
		return mult * lastProbs[lastProbs.length - 1]
				/ (1 + len - lastProbs.length);
	}

	public final float getFirstProb(int si, int len) {
		float nullDistProb = 1f / (1 + len);
		float nonNullDistProb = 1 - nullDistProb;
		float mult = nonNullDistProb;
		if (si >= len) {
			mult = nullDistProb;
			si -= len;
		}
		float norm = 1f;
		if (len < firstProbs.length) {
			float sum = 0f;
			for (int i = 0; i < len; i++) {
				sum += firstProbs[i];
			}
			norm = 1f / sum;
		}
		mult *= norm;
		if (si < firstProbs.length - 1)
			return mult * firstProbs[si];
		return mult * firstProbs[firstProbs.length - 1]
				/ (1 + len - firstProbs.length);
	}

	/**
	 * get the probability of going from sprev to the null word at si in a
	 * sentence of length len -- currently uniform
	 */
	public final float getNullDistProb(int si, int sprev, int len) {
		// this implements a single null -- I think.
		float nullDistProb = 1f / (1 + len);
		if (si == sprev)
			return nullDistProb;
		return 0;
	}

	private final float tooBig(int sprev, int len) {
		int biggestDiff = len - 1 - sprev + diffZero;
		int numBigs = biggestDiff - (distProbs.length - 1);
		return positiveOverflow / numBigs;
	}

	private final float tooSmall(int sprev, int len) {
		int smallestDiff = 0 - sprev + diffZero;
		int numSmalls = -smallestDiff;
		return negativeOverflow / numSmalls;
	}

	// XXX check that this is correctly giving normalization when
	// we have nulls.
	public final float getNorm(int sprev, int len) {
		if (!cache.cachesValid)
			cache.reset();
		if (cache.getNorm(sprev, len) < 0) {
			float sum = 0;
			for (int i = 0; i < len; i++) {
				int diff = i - sprev + diffZero;
				if (diff >= distProbs.length) {
					sum += tooBig(sprev, len);
				} else if (diff < 0) {
					sum += tooSmall(sprev, len);
				} else
					sum += distProbs[diff];
			}
			cache.setNorm(sprev, len, sum);
		}
		return cache.getNorm(sprev, len);
	}

	public final void addToCount(float prob, int sprev, int si) {
		assert !Double.isNaN(prob);
		cache.invalidate();
		int diff = si - sprev + diffZero;
		if (diff >= distProbs.length) {
			positiveOverflow += prob;
		} else if (diff < 0) {
			negativeOverflow += prob;
		} else
			distProbs[diff] += prob;
	}

	public final void addToNullCount(float prob, int sprev, int si) {
	}

	public final void addToFirst(float prob, int si) {
		if (si >= firstProbs.length)
			si = firstProbs.length - 1;
		firstProbs[si] += prob;
	}

	public final void addToLast(float prob, int si, int len) {
		// we index last by len-si
		si = (len - 1) - si;
		if (si >= lastProbs.length)
			si = lastProbs.length - 1;
		lastProbs[si] += prob;
	}

	public void printTable() {
		System.out.println("    Negative Overflow: " + negativeOverflow);
		for (int i = 0; i < distProbs.length; i++) {
			System.out.println("    diff[" + (i - diffZero) + "]: "
					+ distProbs[i]);
		}
		System.out.println("    Positive Overflow: " + positiveOverflow);
		for (int i = 0; i < distProbs.length; i++) {
			System.out.println("    first[" + (i) + "]: " + firstProbs[i]);
		}
		for (int i = 0; i < distProbs.length; i++) {
			System.out.println("    last[-" + (i) + "]: " + lastProbs[i]);
		}

	}

	/**
	 * Cache for a distortion table. This makes the code bigger but speeds
	 * things up quite a bit.
	 */
	private class Cache {
		int cacheSize; // = 200; // should be > length of longest sentence

		public boolean cachesValid = false;

		float[][] normalizations; // length, sprev

		float[][] valueCache; // length, diff

		float[][] negValueCache; // length, diff

		public Cache(int cacheSize) {
			this.cacheSize = cacheSize;
			normalizations = new float[cacheSize][cacheSize];
			valueCache = new float[cacheSize][cacheSize];
			negValueCache = new float[cacheSize][cacheSize];
		}

		public final void reset() {
			cachesValid = true;
			for (int i = 0; i < cacheSize; i++) {
				for (int j = 0; j < cacheSize; j++) {
					normalizations[i][j] = -1;
				}
			}
			for (int i = 0; i < cacheSize; i++) {
				for (int j = 0; j < cacheSize; j++) {
					valueCache[i][j] = -1;
					negValueCache[i][j] = -1;
				}
			}
		}

		public final void invalidate() {
			cachesValid = false;
		}

		public final float getNorm(int sprev, int len) {
			return normalizations[sprev][len];
		}

		public final void setNorm(int sprev, int len, float p) {
			normalizations[sprev][len] = p;
		}

		public final float getVal(int si, int sprev, int len) {
			int diff = si - sprev + diffZero;
			return (diff > 0) ? valueCache[len][diff]
					: negValueCache[len][-diff];
		}

		public final void setVal(int si, int sprev, int len, float p) {
			int diff = si - sprev + diffZero;
			if (diff > 0)
				valueCache[len][diff] = p;
			else
				negValueCache[len][-diff] = p;
		}
	}

}
